/**
** @file:  invert.cpp
** @brief:
**/

#include "invert.h"
#include "operator_mpi.h"
#include <stdlib.h>
#include <math.h>
#include <iostream>
#include <utils.h>
#include <x86intrin.h>
#include "check.h"
using namespace std;

uint64_t count_mpi_time01 = 0;
uint64_t count_mpi_time02 = 0;
uint64_t count_mpi_time03 = 0;
int onlyonce_01 = 3;
int onlyonce_02 = 20;
int onlyonce_03 = 20;
int global_rank;

int CGinvert(complex<double> *src_p, complex<double> *dest_p, complex<double> *gauge[4],
             const double mass, const int max, const double accuracy, int *subgs, int *site_vec)
{
    lattice_gauge U(gauge, subgs, site_vec);
    lattice_fermion src(src_p, subgs, site_vec);
    lattice_fermion dest(dest_p, subgs, site_vec);
    CGinvert(src, dest, U, mass, max, accuracy);
    return 0;
}

int CGinvert(
    lattice_fermion &src,   // M^\dagger M * dest = M^\dagger src, src means 'b'
    lattice_fermion &dest,  // M^\dagger M * dest = M^\dagger src, dest means 'x'
    lattice_gauge &U,       // 组态
    const double mass,      // 夸克质量
    const int max,          // 默认给定值迭代次数 500
    const double accuracy   // 默认精度 1.0e-9
    )
{
    int myrank;
    MPI_Comm_rank(MPI_COMM_WORLD, &myrank);
    global_rank = myrank;
    lattice_fermion r0(src.subgs, src.site_vec);    // 残差, 同时也是一开始的梯度
    lattice_fermion r1(src.subgs, src.site_vec);    // 残差, 同时也是一开始的梯度
    lattice_fermion q(src.subgs, src.site_vec);     // q = M^T qq = M^T M p
    lattice_fermion qq(src.subgs, src.site_vec);    // qq = M p
    lattice_fermion p(src.subgs, src.site_vec);     // 共轭向量
    lattice_fermion Mdb(src.subgs, src.site_vec);   // M^\dagger src
    lattice_fermion tmp(src.subgs, src.site_vec);

    complex<double> aphi(0);
    complex<double> beta(0);
    // M = L(M~)U
    Dslash_left(src, tmp, U, mass, false);      // L^{-1} src = b'
    Dslash_middle(tmp, Mdb, U, mass, true);     // (M~)^\dagger b'
    // Dslash(src, Mdb, U, mass, true);            // M^\dagger src
    for (int i = 0; i < dest.size; i++) {
        // dest.A[i] = 1.0 * rand() / RAND_MAX;    // 随机一个初始解 x
        Mdb.A[i] / ((4.0 + mass) * (4.0 + mass));
    }
    // dest.A[0] = 1.0;
#ifdef DEBUG
    double src_nr2 = norm_2(src);
    double dest_nr2 = norm_2(dest);
    if (myrank == 0) {
        cout << "|src|^2 = " << src_nr2 << " , |dest|^2 = " << dest_nr2 << endl;
    }
#endif
    // Dslash(dest, tmp, U, mass, false);          // tmp = M x
    // Dslash(tmp, r0, U, mass, true);             // r0 = M^\dagger M x
    Dslash_middle(dest, tmp, U, mass, false);
    Dslash_middle(tmp, r0, U, mass, true);

    for (int i = Mdb.size/2; i < Mdb.size; i++) {
        r0.A[i] = Mdb.A[i] - r0.A[i];           // r0 = M^\dagger src - M^\dagger M * dest
    }

    int count_error_times = 0;
    for (int f = 1; f < max; f++) {
        count_mpi_time01 = 0;
        count_mpi_time02 = 0;
        count_mpi_time03 = 0;
        auto startTime = Clock::now();
        if (f == 1) {
            for (int i = 0; i < r0.size; i++)
                p.A[i] = r0.A[i];                   // 一开始搜索方向即为负梯度方向，即为残差向量方向
        } else {
            beta = vector_p_new(r0, r0) / vector_p_new(r1, r1);
            double beta_real = beta.real();
            double* pA = (double*)p.A;
            double* r0A = (double*)r0.A;
            const int max = r0.size*2;
            for (int i = max/2; i < max; i++)
                pA[i] = r0A[i] + beta_real * pA[i];
                // p.A[i] = r0.A[i] + beta * p.A[i];   // 之后共轭向量为 残差 和 上一次共轭向量的线性组合
        }
        auto endTime = Clock::now();
        auto compTime0 = chrono::duration_cast<chrono::microseconds>(endTime - startTime);
        startTime = Clock::now();
        // Dslash(p, qq, U, mass, false);  // qq = M p, A = M^T M
        Dslash_middle(p, qq, U, mass, false);  // qq = M p, A = M^T M
        endTime = Clock::now();
        auto compTime1 = chrono::duration_cast<chrono::microseconds>(endTime - startTime);
        startTime = Clock::now();
        // Dslash(qq, q, U, mass, true);   // q = M^T p, A = M^T M
        Dslash_middle(qq, q, U, mass, true);
        endTime = Clock::now();
        auto compTime2 = chrono::duration_cast<chrono::microseconds>(endTime - startTime);
        startTime = Clock::now();

        aphi = vector_p_new(r0, r0) / vector_p_new(qq, qq);   // \frac{r^T r}{p^T A p}
        // aphi = vector_p(r0, r0) / vector_p(p, q);   // \frac{r^T r}{p^T A p}
        endTime = Clock::now();
        auto compTime3 = chrono::duration_cast<chrono::microseconds>(endTime - startTime);
        startTime = Clock::now();

        for (int i = dest.size/2; i < dest.size; i++){
            dest.A[i] = dest.A[i] + aphi * p.A[i];  // x = x + \alpha p
        }
        // 到时候来个替换，一次 r0，一次 r1，就不用拷贝了
        for (int i = r1.size/2; i < r1.size; i++)
            r1.A[i] = r0.A[i];
        double aphi_real = -aphi.real();
        // double* r0_double = (double*)(r0.A);
        // double* q_double = (double*)(q.A);
        for (int i = r0.size/2; i < r0.size; i++){
            // r0_double[i] = r0_double[i] - aphi_real * q_double[i];
            // r0_double[i+1] = r0_double[i+1] - aphi_real * q_double[i+1];
            // r0_double[i+2] = r0_double[i+2] - aphi_real * q_double[i+2];
            // r0_double[i+3] = r0_double[i+3] - aphi_real * q_double[i+3];
            r0.A[i] = r0.A[i] - aphi * q.A[i];
        }

        endTime = Clock::now();
        auto compTime4 = chrono::duration_cast<chrono::microseconds>(endTime - startTime);
        startTime = Clock::now();
        double rsd2 = norm_2_half(r0);   // 求复向量内积 sum(conj(a[i]) * a[i])
        double rsd = sqrt(rsd2);    // 复向量的模

        endTime = Clock::now();
        auto compTime5 = chrono::duration_cast<chrono::microseconds>(endTime - startTime);
        startTime = Clock::now();

        if (rsd < accuracy) {
            for (int i = 0; i < dest.size; i++){
                // dest.A[i] = dest.A[i] + aphi * p.A[i];  // x = x + \alpha p
                tmp.A[i] = dest.A[i] + aphi * p.A[i];  // x = x + \alpha p
            }
            for(int i=0; i<dest.size/2; i++){
                tmp.A[i] = Mdb.A[i] / ((4.0 + mass) * (4.0 + mass));
            }
            Dslash_right(tmp, dest, U, mass, false);
            // 计算真正误差，因为误差要放大，不满足则迭代回去。当然，可能实在满足不了
            // double Mx_b_r, MMx_Mb_r;
            // check_inv(src, tmp, U, mass, Mx_b_r, MMx_Mb_r);
            // if(Mx_b_r < accuracy && MMx_Mb_r < accuracy){
            // if((Mx_b_r < accuracy && MMx_Mb_r < accuracy) || count_error_times >= 10){
                if (myrank == 0) {
                    cout << "CG: " << f << " iterations, convergence residual |r| = " << rsd << endl;
                }
                break;
            // }
            // else{
            //     count_error_times++;
            //     if (myrank == 0)
            //         cout << "Rechecking due to the propagation of error." << endl;
            // }
        }
#ifndef VERBOSE_SIMPLE
        if (myrank == 0) {
            cout << "CG: " << f << " iter, rsd |r| = " << rsd << " \t";
            printf("Computing time(ms): %ld %ld %ld %ld %ld %ld \t", 
                                                    compTime0.count(), compTime1.count(),
                                                    compTime2.count(), compTime3.count(), 
                                                    compTime4.count(), compTime5.count());
            printf("MPI com time(ms): %ld %ld %ld\n", count_mpi_time01, count_mpi_time02, count_mpi_time03);
        }
#endif
    }
    return 0;
}

// https://zhuanlan.zhihu.com/p/369204956
// https://blog.csdn.net/seventonight/article/details/115863380
// https://max.book118.com/html/2018/0608/171404230.shtm
int BiCGSTABinvert(
    lattice_fermion &src,   // M^\dagger M * dest = M^\dagger src, src means 'b'
    lattice_fermion &dest,  // M^\dagger M * dest = M^\dagger src, dest means 'x'
    lattice_gauge &U,       // 组态
    const double mass,      // 夸克质量
    const int max,          // 默认给定值迭代次数 500
    const double accuracy   // 默认精度 1.0e-9
    )
{
    int myrank;
    MPI_Comm_rank(MPI_COMM_WORLD, &myrank);
    global_rank = myrank;
    lattice_fermion r0(src.subgs, src.site_vec);
    lattice_fermion r(src.subgs, src.site_vec);
    lattice_fermion p(src.subgs, src.site_vec);
    lattice_fermion v(src.subgs, src.site_vec);
    lattice_fermion t(src.subgs, src.site_vec);
    lattice_fermion Lb(src.subgs, src.site_vec);
    // lattice_fermion s(src.subgs, src.site_vec);
    lattice_fermion tmp(src.subgs, src.site_vec);
    complex<double> rho(1), alpha(1), omega(1), beta(1), rho_old(1);
    Dslash_left(src, Lb, U, mass, false);       // L^{-1} src = b'
    for (int i = 0; i < dest.size; i++) {
        // dest.A[i] = 1.0 * rand() / RAND_MAX;
        dest.A[i] = Lb.A[i] / (4.0 + mass);
    }
    Dslash_middle(dest, r0, U, mass, false);    // r0 = Mx
    // Dslash(dest, r0, U, mass, false);
    for (int i = r.size/2; i < r.size; i++) {
        r0.A[i] = Lb.A[i] - r0.A[i];            // r0 = Mx - Lb
    }
    for (int i = r.size/2; i < r.size; i++) {
        r.A[i] = r0.A[i];                       // r = r0 = Mx - Lb
    }
    for(int iterf = 1; iterf < max; iterf++){
        rho = vector_p_half(r0, r);
        beta = (rho / rho_old) * (alpha / omega);
        rho_old = rho;
        if(iterf == 1){
            for (int i = r.size/2; i < r.size; i++)
                p.A[i] = r0.A[i];                           // p = Lb - Mx
        }
        else{
            for (int i = r.size/2; i < r.size; i++)
                p.A[i] = r.A[i] + beta * (p.A[i] - omega * v.A[i]);
        }
        // Dslash(p, v, U, mass, false); // v = Mp
        Dslash_middle(p, v, U, mass, false);                // v = M p
        alpha = rho / vector_p_half(r0, v);                 // alpha = rho / (r0, v)
        for (int i = r.size/2; i < r.size; i++)             // s = r - alpha v, s 用 r 暂存即可
            r.A[i] = r.A[i] - alpha * v.A[i];
        // Dslash(s, t, U, mass, false); // t = M s
        Dslash_middle(r, t, U, mass, false);                // t = M s
        omega = vector_p_half(t, r) / vector_p_half(t, t);  // omega = (t, s) / (t, )
        for (int i = r.size/2; i < r.size; i++)             // x = x + alpha p + omega s
            dest.A[i] = dest.A[i] + alpha * p.A[i] + omega * r.A[i];
        for (int i = r.size/2; i < r.size; i++)             // r = s - omega t
            r.A[i] = r.A[i] - omega * t.A[i];
        double rsd2 = norm_2_half(r);
        double rsd = sqrt(rsd2);
        if(rsd < accuracy){
            for(int i=0; i<dest.size/2; i++){
                dest.A[i] = Lb.A[i] / (4.0 + mass);
            }
            Dslash_right(dest, tmp, U, mass, false);
            for (int i = 0; i < r.size; i++){
                dest.A[i] = tmp.A[i];
            }
            if (myrank == 0) {
                cout << "CG: " << iterf << " iterations, convergence residual |r| = " << rsd << endl;
            }
            break;
        }
#ifndef VERBOSE_SIMPLE
        if (myrank == 0) {
            cout << "CG: " << iterf << " iter, rsd |r| = " << rsd << endl;
        }
#endif
    }
    return 0;
}

void Dslash(lattice_fermion &src, lattice_fermion &dest, lattice_gauge &U, const double mass,
            const bool dagger)
{
    dest.clean();
    lattice_fermion tmp(src.subgs, src.site_vec);
    DslashEE(src, tmp, mass);
    dest = dest + tmp;
    DslashOO(src, tmp, mass);
    dest = dest + tmp;
    Dslashoffd(src, tmp, U, dagger, 0); // cb=0, EO
    dest = dest + tmp;
    Dslashoffd(src, tmp, U, dagger, 1);
    dest = dest + tmp;
}

void Dslash_new(
    lattice_fermion &src,   // 相当于 b
    lattice_fermion &dest,  // 相当于 x
    lattice_gauge &U,       // 组态
    const double mass,      // 夸克质量
    const bool dagger       // 是否要取 M^{dagger}
)
{
    dest.clean();
    lattice_fermion tmp(src.subgs, src.site_vec);
    // DslashEE(src, tmp, mass);   // tmp[0:V/2] = (4.0 + mass) * src[0:V/2]
    // dest = dest + tmp;
    // DslashOO(src, tmp, mass);   // tmp[V/2: ] = (4.0 + mass) * src[V/2: ]
    // dest = dest + tmp;

    const double a = 4.0;
    int subgrid_vol = (src.subgs[0] * src.subgs[1] * src.subgs[2] * src.subgs[3]);
    for (int i = 0; i < subgrid_vol * 3 * 4; i++) {
        dest.A[i] = (a + mass) * src.A[i];
    }

    // Warning !!! a+mass 写到了 Dslashoffd 中间部分 t 遍历部分
    Dslashoffd_new(src, dest, U, dagger, 0); // cb=0, EO
    // dest = dest + tmp;
    Dslashoffd_new(src, dest, U, dagger, 1);
    // dest = dest + tmp;
}


void Dslash_left(
    lattice_fermion &src,   // 相当于 b
    lattice_fermion &dest,  // 相当于 x
    lattice_gauge &U,       // 组态
    const double mass,      // 夸克质量
    const bool dagger       // 是否要取 M^{dagger}
)
{
    dest.clean();
    Dslashoffd_new(src, dest, U, false, 1); // cb=0, EO
    const double a = 4.0;
    int subgrid_vol = (src.subgs[0] * src.subgs[1] * src.subgs[2] * src.subgs[3]);
    int subgrid_vol_cb = (subgrid_vol) >> 1;
    for (int i = 0; i < subgrid_vol_cb * 3 * 4; i++) {
        dest.A[i] = src.A[i];
    }
    for (int i = subgrid_vol_cb * 3 * 4; i < subgrid_vol * 3 * 4; i++) {
        dest.A[i] = src.A[i] - (1/(a+mass)) * dest.A[i];
    }
}

void Dslash_left_forward(
    lattice_fermion &src,   // 相当于 b
    lattice_fermion &dest,  // 相当于 x
    lattice_gauge &U,       // 组态
    const double mass,      // 夸克质量
    const bool dagger       // 是否要取 M^{dagger}
)
{
    dest.clean();
    const double a = 4.0;
    int subgrid_vol = (src.subgs[0] * src.subgs[1] * src.subgs[2] * src.subgs[3]);
    int subgrid_vol_cb = (subgrid_vol) >> 1;
    if(dagger){
        Dslashoffd_new(src, dest, U, true, 0); // cb=0, EO
        for (int i = 0; i < subgrid_vol_cb * 3 * 4; i++) {
            dest.A[i] = src.A[i] + (1/(a+mass)) * dest.A[i];
        }
        for (int i = subgrid_vol_cb * 3 * 4; i < subgrid_vol * 3 * 4; i++) {
            dest.A[i] = src.A[i];
        }
    }
    else{
        Dslashoffd_new(src, dest, U, false, 1); // cb=0, EO
        for (int i = 0; i < subgrid_vol_cb * 3 * 4; i++) {
            dest.A[i] = src.A[i];
        }
        for (int i = subgrid_vol_cb * 3 * 4; i < subgrid_vol * 3 * 4; i++) {
            dest.A[i] = src.A[i] + (1/(a+mass)) * dest.A[i];
        }
    }
}

void Dslash_right(
    lattice_fermion &src,   // 相当于 b
    lattice_fermion &dest,  // 相当于 x
    lattice_gauge &U,       // 组态
    const double mass,      // 夸克质量
    const bool dagger       // 是否要取 M^{dagger}
)
{
    dest.clean();
    Dslashoffd_new(src, dest, U, false, 0); // cb=0, EO
    const double a = 4.0;
    int subgrid_vol = (src.subgs[0] * src.subgs[1] * src.subgs[2] * src.subgs[3]);
    int subgrid_vol_cb = (subgrid_vol) >> 1;
    for (int i = 0; i < subgrid_vol_cb * 3 * 4; i++) {
        dest.A[i] = src.A[i] - (1/(a+mass)) * dest.A[i];
    }
    for (int i = subgrid_vol_cb * 3 * 4; i < subgrid_vol * 3 * 4; i++) {
        dest.A[i] = src.A[i];
    }
}

void Dslash_right_forward(
    lattice_fermion &src,   // 相当于 b
    lattice_fermion &dest,  // 相当于 x
    lattice_gauge &U,       // 组态
    const double mass,      // 夸克质量
    const bool dagger       // 是否要取 M^{dagger}
)
{
    dest.clean();
    const double a = 4.0;
    int subgrid_vol = (src.subgs[0] * src.subgs[1] * src.subgs[2] * src.subgs[3]);
    int subgrid_vol_cb = (subgrid_vol) >> 1;
    if(dagger){
        Dslashoffd_new(src, dest, U, true, 1); // cb=0, EO
        for (int i = 0; i < subgrid_vol_cb * 3 * 4; i++) {
            dest.A[i] = src.A[i];
        }
        for (int i = subgrid_vol_cb * 3 * 4; i < subgrid_vol * 3 * 4; i++) {
            dest.A[i] = src.A[i] + (1/(a+mass)) * dest.A[i];
        }
    } else {
        Dslashoffd_new(src, dest, U, false, 0); // cb=0, EO
        for (int i = 0; i < subgrid_vol_cb * 3 * 4; i++) {
            dest.A[i] = src.A[i] + (1/(a+mass)) * dest.A[i];
        }
        for (int i = subgrid_vol_cb * 3 * 4; i < subgrid_vol * 3 * 4; i++) {
            dest.A[i] = src.A[i];
        }
    }
}

void Dslash_middle_origin(
    lattice_fermion &src,   // 相当于 b
    lattice_fermion &dest,  // 相当于 x
    lattice_gauge &U,       // 组态
    const double mass,      // 夸克质量
    const bool dagger       // 是否要取 M^{dagger}
)
{
    dest.clean();
    Dslashoffd_new(src, dest, U, dagger, 0); // tmp[0:V/2] = Meo * src[V/2:V]
    Dslashoffd_new(dest, dest, U, dagger, 1); // tmp2[V/2:V] = Moe * tmp[0:V/2]
    const double a = 4.0;
    int subgrid_vol = (src.subgs[0] * src.subgs[1] * src.subgs[2] * src.subgs[3]);
    int subgrid_vol_cb = (subgrid_vol) >> 1;
    for (int i = 0; i < subgrid_vol_cb * 3 * 4; i++) {
        dest.A[i] = (a + mass) * src.A[i];
    }
    for (int i = subgrid_vol_cb * 3 * 4; i < subgrid_vol * 3 * 4; i++) {
        dest.A[i] = (a + mass) * src.A[i] - (1/(a+mass)) * dest.A[i];
    }
}

void Dslash_middle(
    lattice_fermion &src,   // 相当于 b
    lattice_fermion &dest,  // 相当于 x
    lattice_gauge &U,       // 组态
    const double mass,      // 夸克质量
    const bool dagger       // 是否要取 M^{dagger}
)
{
    dest.clean();
    Dslashoffd_new(src, dest, U, dagger, 0); // tmp[0:V/2] = Meo * src[V/2:V]
    Dslashoffd_new(dest, dest, U, dagger, 1); // tmp2[V/2:V] = Moe * tmp[0:V/2]
    const double a = 4.0;
    int subgrid_vol = (src.subgs[0] * src.subgs[1] * src.subgs[2] * src.subgs[3]);
    int subgrid_vol_cb = (subgrid_vol) >> 1;
    // for (int i = 0; i < subgrid_vol_cb * 3 * 4; i++) {
    //     dest.A[i] = (a + mass) * src.A[i];
    // }
    for (int i = subgrid_vol_cb * 3 * 4; i < subgrid_vol * 3 * 4; i++) {
        dest.A[i] = (a + mass) * src.A[i] - (1/(a+mass)) * dest.A[i];
    }
}